import tensorflow as tf
import pathlib
from tensorflow import keras


def getData():
    # 加载数据集
    path = "../../static/img"
    # 解析目录
    data_dir = pathlib.Path(path)

    # keras 加载数据集
    batch_size = 32
    img_height = 180
    img_width = 180

    # 使用 80% 的图像进行训练，20% 的图像进行验证。
    class_names = ['Battery', 'BrickAndTileCeramics', 'Cans', 'cigarette',
                   'Fruits', 'NO_RUBBISH', 'Vegetables', 'WaterBottle']
    train_ds = keras.utils.image_dataset_from_directory(
        data_dir,
        validation_split=0.2,
        subset="training",
        image_size=(img_height, img_width),
        batch_size=batch_size,
        shuffle=True,
        seed=123,
        interpolation='bilinear',
        crop_to_aspect_ratio=True,
        labels='inferred',
        class_names=class_names,
        color_mode='rgb'
    )

    val_ds = keras.utils.image_dataset_from_directory(
        data_dir,
        validation_split=0.2,
        subset="validation",
        image_size=(img_height, img_width),
        batch_size=batch_size,
        shuffle=True,
        seed=123,
        interpolation='bilinear',
        crop_to_aspect_ratio=True,
        labels='inferred',
        class_names=class_names,
        color_mode='rgb'
    )

    AUTOTUNE = tf.data.AUTOTUNE
    train_ds = train_ds.cache().prefetch(buffer_size=AUTOTUNE)
    val_ds = val_ds.cache().prefetch(buffer_size=AUTOTUNE)
    return train_ds, val_ds, len(class_names), img_width, img_height
